# #!/usr/bin/env python
# # -*- coding: utf-8 -*-
# # @Copyright © 启明星辰 版权所有
# # @Time    : 2024/3/27 18:00
# # @Author  : 胡浩浩
# # @File    : mitmproxy_plugin.py
# # @IDE     : PyCharm
# # @description :
import copy
import json
import re

from bs4 import BeautifulSoup
from mitmproxy import http, ctx
import base64
from collections import defaultdict
from io import BytesIO
from typing import Any
from multipart import parse_form as orig_parse_form
from pydantic import BaseModel
import urllib.parse
from multipart.multipart import Field, File
from ruamel import yaml
from urllib.parse import parse_qs, parse_qsl


class FileModel(BaseModel):
    filename: str
    base64: str


def unpack_single_item_lists(data: dict[Any, list[Any]]):
    """
    如果dict中某个键的值是单元素list，则直接拆开
    eg: {"a": [1]} -> {"a": 1}
    """
    return {
        k: (v if isinstance(v, list) and len(v) > 1 else v[0]) for k, v in data.items()
    }


def parse_form(
        headers: dict[str, str],
        stream: BytesIO,
        chunk_size: int = 1048576,
) -> dict[str, list[Any]]:
    items: dict[str, list[Any]] = defaultdict(list)

    def on_field(field: Field):
        if not (isinstance(field.field_name, bytes) and isinstance(field.value, bytes)):
            raise TypeError("什么，读取到的值竟然不是 bytes？")  # impossible
        name = field.field_name.decode()
        value = field.value.decode()
        items[name].append(value)

    def on_file(file: File):
        if not isinstance(file.field_name, bytes):
            raise TypeError("什么，读取到的值竟然不是bytes？")  # impossible
        name = file.field_name.decode()
        # 需要将文件offset重置为0，否则无内容可读
        file.file_object.seek(0)
        value = FileModel(
            filename=file.file_name,
            base64=base64.b64encode(file.file_object.read()).decode(),
        ).model_dump()
        items[name].append(value)

    orig_parse_form(
        headers=headers,
        input_stream=stream,
        on_field=on_field,
        on_file=on_file,
        chunk_size=chunk_size,
    )

    return items


def item_to_dict(obj):
    """object to dict"""
    return dict(obj.items()) if obj.items() else None


def generate_jmespath(node, path='', expressions=None):
    """
    Generate JMESPath expressions for each level of a nested dictionary.

    Args:
    - node (dict): The nested dictionary.
    - path (str): The current path.
    - expressions (dict): A dictionary to store JMESPath expressions and values.

    Returns:
    - dict: A dictionary where keys are JMESPath expressions and values are corresponding values.
    """
    if expressions is None:
        expressions = {}
    if isinstance(node, dict):
        for key, value in node.items():
            new_path = f"{path}.{key}" if path else key
            if isinstance(value, dict):
                generate_jmespath(value, new_path, expressions)
                continue
            else:
                if isinstance(value, str) or isinstance(value, int):
                    expressions[new_path] = value
    else:
        pass
    return expressions


def jmespath_dict_to_str_list(jmespath_dict: dict):
    """
    Generate jmespath dict to str list
    :param jmespath_dict:
    :return:
    """
    express_list = []
    if isinstance(jmespath_dict, dict):
        for key, value in jmespath_dict.items():
            # print(f"{key}==`{value}`")
            if isinstance(value, int):
                if str(value) == "True" or str(value) == "False":
                    continue
                express_list.append(f'{key}==`{value}`')
            elif isinstance(value, dict):
                express_list.append(f'{key}=={value}')
            else:
                express_list.append(f'{key}==`"{value}"`')
    if not express_list:
        express_list = None
    return express_list


def re_express_parse(html_content):
    # 使用Beautiful Soup解析HTML
    soup = BeautifulSoup(html_content, 'html.parser')

    # 获取title标签中的文本数据
    title = soup.title.get_text()

    # 打印标题数据
    # print("Title-------------:", title)
    return title


product_list = ["csa", "asm", "cdn", "ac", "cl", "sa", "ios", 'aism', 'xdr', 'isop-soc', 'ics-csa',
                'leadsec-rs', 'mdc', "gpt", 'sdc', "cssp", "isop", "csa-cmsi", "panguip-portal",
                "panguip-bridge", "panguip-cypher", "panguip-trust", "panguip-console", "panguip-ti",
                "panguip-sno", "panguip-mdc", "panguip-range", "panguip-kaas", "panguip-gpt","usm","xfc"]


# host_list = ["192.", "10."]
# port_list = [8081]

# path_end = ["js", "png", "jpg", "json", "css"]


def check_product(product: str) -> bool:
    """
        Verify that the product name is in the URL, return false if it is not, and true if it is
    :param product:
    :return: true or false
    """
    flag = False

    for h in product_list:
        if h in product:
            return True
    return flag

def find_parent_key(data, target_key):
    """
    在嵌套字典中找到目标键的所有父节点的键名。

    :param data: 目标字典
    :param target_key: 目标键名
    :return: 包含所有父节点键名的列表，如果未找到则返回空列表
    """
    data_list = []

    def search(parent, target):
        if isinstance(parent, dict):
            for key, value in parent.items():
                if isinstance(value, dict):
                    if target in value:
                        data_list.append(key)
                    # 递归搜索子字典
                    search(value, target)
        return

    search(data, target_key)
    return data_list


class RewriteContent:

    def __init__(self):
        # 设置 Mitmweb 相关选项
        self.file_suffix_list = [
            ".json$",
            ".svg$",
            ".jpg$",
            ".css$",
            ".js$",
            ".png$",
            ".icon$",
            ".ttf$",
            ".tif$",
            ".bmp$",
            ".gif$",
            ".pcx$",
            ".tga$",
            ".exif$",
            ".fpx$",
            ".psd$",
            ".cdr$",
            ".pcd$",
            ".dxf$",
            ".ufo$",
            ".eps$",
            ".ai$",
            ".raw$",
            ".WMF$",
            ".webp$",
            ".avif$",
            ".apng$",
        ]
        self.options = {
            "view_filter": ' & '.join(
                ['!' + suffix for suffix in self.file_suffix_list])
        }

    def running(self):
        # 设置选项
        for option, value in self.options.items():
            ctx.options.__setattr__(option, value)
            ctx.log.info(f"Set {option} to {value}")

    def response(self, flow: http.HTTPFlow):
        request = flow.request
        response = flow.response
        # print(request.pretty_host)
        # print(request.port)
        # print(type(request.port))

        query_params = item_to_dict(request.query)
        headers = item_to_dict(request.headers)
        mime_type = headers.get("content-type").split(";")[0] if headers.get("content-type") else \
            headers.get("Content-Type").split(";")[0] if headers.get("Content-Type") else None
        referer = headers.get("referer") if headers.get("referer") else headers.get("Referer")
        if referer:
            referer = urllib.parse.unquote(referer)
        url = request.url
        host = re.findall(r"(http.*//.*?)/", url).pop() if re.findall(r"(http.*//.*)/", url) else None
        path = request.path.split("?")[0]
        websocket_origin = ""
        if flow.websocket:
            websocket_origin = "wss://${ip}:${port}"
            print("我拿到的websocket消息", flow.websocket.messages)
        product = ""
        url_product = path.split("/")[1] if len(path.split("/")) >= 2 else ""
        if not check_product(url_product):
            return
        if url_product in product_list:
            product = url_product
            path = "/${product}" + path[len(product) + 1:]
        if referer:
            referer = referer.replace(f"{host}/{product}", "${origin}/${product}") if product else referer.replace(
                f"{host}", "${origin}")
        if product:
            step_name = path[11:].replace("/", "_").replace("-", "_").replace(".", "_")
        else:
            step_name = path.replace("/", "_").replace("-", "_").replace(".", "_")
        method = request.method
        if mime_type and "multipart/form-data" in mime_type:
            boundary = request.content.decode("latin-1").split("\r\n")[0]
            print(boundary)
            data = unpack_single_item_lists(
                parse_form(
                    headers={
                        "Content-Type": f"multipart/form-data; boundary={boundary[2:]}"
                    },
                    stream=BytesIO(request.content),
                )
            )
            # print(json.dumps(data))
            # if b"Content-Type: application/json" in request.content:
            parent_key = find_parent_key(data, 'base64')
        #
            if parent_key:
                for i in parent_key:
                    try:
                        base64_data = data.get(i).get("base64")
                        print(base64_data)
                        print(json.loads(base64.b64decode(base64_data).decode("utf-8")))
                        del data[i]["base64"]
                        data[i]["content"] = json.loads(base64.b64decode(base64_data).decode("utf-8"))
                    except:
                        print("啥也不用干")

        elif mime_type and "x-www-form-urlencoded" in mime_type:
            # 将键值对参数转化为JSON
            # 解码字节字符串
            decoded_string = request.content.decode('utf-8')
            # 切分键值对
            key_value_pairs = decoded_string.split('&')
            # 创建空字典
            params_dict = {}
            # 解析每个键值对，并添加到字典中
            for pair in key_value_pairs:
                if pair and "=" in pair:
                    print("huhaohao",pair)
                    key, value = pair.split('=')
                    # URL 解码键和值
                    key = urllib.parse.unquote(key.replace("+", " "))
                    value = urllib.parse.unquote(value.replace("+", " "))
                    if params_dict.get(key) and isinstance(params_dict.get(key), str):
                        source_value = params_dict.get(key)
                        params_dict[key] = []
                        params_dict[key].append(source_value)
                        params_dict[key].append(value)
                    elif params_dict.get(key) and isinstance(params_dict.get(key), list):
                        params_dict[key].append(value)
                    else:
                        params_dict[key] = value
                else:
                    params_dict[pair]=None
            data = params_dict
        else:
            print("我定位到原因了-----------start")
            print(request.content)
            print("我定位到原因了-----------end")
            try:
                data = json.loads(request.content.decode('utf-8')) if request.content else None
            except:
                print("不是标准的json")
                data= request.content.decode('utf-8')
        response_content_type = response.headers.get("content-type") or response.headers.get("Content-Type")
        express_list = []
        checker_regex = []
        if response_content_type and "application/json" in response_content_type:
            if response.get_content():
                expressions = generate_jmespath(response.json())
                express_list = jmespath_dict_to_str_list(expressions)
        else:
            if response and response_content_type and "text/html" in response_content_type:
                title = re_express_parse(response.content)
                checker_regex.append({
                    "pattern": f"{title}",
                    "description": f"验证<title>标签里面的值为{title}"
                })
        yml_case_temp = {
            "suite_name": "场景名称_自定义",
            "suite_desc": "场景描述信息_自定义",
            "variables": None,
            "test_steps": [
                {
                    "step_name": step_name if step_name else "步骤名称_自定义",
                    "step_desc": "步骤描述信息_自定义",
                    "tags": None,
                    "weight": 1,
                    "rerun": None,
                    "req_info": {
                        "origin": websocket_origin if websocket_origin else "${origin}",
                        "method": method,
                        "path": path,
                        "mime_type": mime_type,
                        "headers": {
                            "Referer": referer
                        }
                    },
                    "test_cases": [
                        {
                            "case_desc": "用例描述_自定义",
                            "timeout": 60,
                            "setup": None,
                            "precondition": None,
                            "use_user": None,
                            "sleep": 0.0,
                            "params": query_params if query_params else None,
                            "data": data if data else None,
                            "save_vars": None,
                            "checker": {
                                "code": 200,
                                "regex": checker_regex if checker_regex else None,
                                "json": express_list if express_list else None,
                                "sql": None
                            },
                            "teardown": None
                        }
                    ]
                }
            ]
        }
        yaml_response = copy.deepcopy(flow.response)
        # print("HHH",flow.response.content)
        filename = "test_" + flow.request.path.split("?")[0].replace("/", "_").replace("-", "_") + ".yml"
        cd = f"attachment; filename={filename}"
        yaml_response.headers["Content-Disposition"] = cd
        yaml_response.headers["content-type"] = "application/json; charset=utf-8"
        yaml_response.headers["content-length"] = str(
            len(yaml_response.raw_content)) if yaml_response.raw_content is not None else None
        yaml_response.set_content(
            bytes(yaml.round_trip_dump(yml_case_temp, default_flow_style=False, allow_unicode=True), "utf-8"))
        flow.__setattr__("yaml_response", yaml_response)


addons = [RewriteContent()]

# 本地启动
#  mitmweb -s .\mitmproxy_plugin.py --set ssl_insecure=true --web-host 192.168.114.124 --web-port 8081 --listen-host 0.0.0.0 --listen-port 8180 --set view_filter='!.css & !.js & !.png & !.icon & !.ttf'

# 中文文档 https://ptorch.com/docs/10/mitmproxy_introduction

# docker启动
# docker run -d -it -p 8180:8180 -p 8081:8081 ptesting_mitmproxy mitmweb -s mitmproxy_plugin.py --set ssl_insecure=true --listen-host 0.0.0.0 --listen-port 8180 --web-host 0.0.0.0 --web-port 8081 --set view_filter='!.css & !.js & !.png & !.icon & !.ttf & !.svg & !.jpg'
